import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset
from tqdm import tqdm
import matplotlib.pyplot as plt
from torch.utils.tensorboard import SummaryWriter  # 导入 TensorBoard
from lstm_dataloader_slide_window import lstm_dataloader_slide_window
from Config import Config

# 加载数据
X_train_tensor, X_test_tensor, y_train_tensor, y_test_tensor = lstm_dataloader_slide_window(window_size=15,
                                                                                            test_size=0.05)
train_data = TensorDataset(X_train_tensor, y_train_tensor)
test_data = TensorDataset(X_test_tensor, y_test_tensor)

config = Config()

# 将数据加载成迭代器
train_loader = torch.utils.data.DataLoader(train_data, batch_size=config.batch_size, shuffle=False)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=config.batch_size, shuffle=False)


# 定义 LSTM 模型
class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=1, dropout=0.2):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        out = self.fc(lstm_out)
        return out


# 超参数
input_size = X_train_tensor.shape[2]
hidden_size = Config.hidden_size
output_size = y_train_tensor.shape[2]
num_layers = Config.num_layers
num_epochs = Config.epochs
learning_rate = Config.learning_rate
dropout = Config.dropout

# 模型、损失函数、优化器
model = LSTMModel(input_size, hidden_size, output_size, num_layers, dropout)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 初始化 TensorBoard
writer = SummaryWriter(log_dir='./runs/lstm_experiment')  # 日志保存路径

# 训练循环
train_losses = []
for epoch in range(num_epochs):
    model.train()
    epoch_loss_sum = 0.0
    train_bar = tqdm(train_loader)
    for batch_idx, data in enumerate(train_bar):
        X_train_batch, y_train_batch = data
        optimizer.zero_grad()
        y_train_pred = model(X_train_batch)
        train_loss = criterion(y_train_pred.view(-1), y_train_batch.view(-1))
        epoch_loss_sum += train_loss.detach().numpy()
        train_loss.backward()
        optimizer.step()
        train_bar.desc = "train epoch[{} / {}] loss:{:.3f}".format(epoch + 1, num_epochs, train_loss)

        # 记录每个 batch 的损失到 TensorBoard
        writer.add_scalar('Train Loss (per batch)', train_loss.item(), epoch * len(train_loader) + batch_idx)

    # 记录每个 epoch 的平均损失到 TensorBoard
    avg_epoch_loss = epoch_loss_sum / len(train_loader)
    writer.add_scalar('Train Loss (per epoch)', avg_epoch_loss, epoch)
    train_losses.append(avg_epoch_loss)

# 保存模型
torch.save(model.state_dict(), './lstm_model.pth')

# 测试模型
model.eval()
accuracy_threshold = 0.1
with torch.no_grad():
    predictions = model(X_test_tensor)
    mse_loss = criterion(predictions, y_test_tensor).item()
    print(f"Mean Squared Error on Test Set: {mse_loss:.4f}")

    true_values = y_test_tensor.numpy()
    predicted_values = predictions.numpy()

    # 计算准确率
    accurate_predictions = np.abs(predicted_values - true_values) < accuracy_threshold * np.abs(true_values)
    accuracy_percentage = np.mean(accurate_predictions) * 100
    print(f"Accuracy within {accuracy_threshold * 100}% of true values: {accuracy_percentage:.2f}%")

    # 记录测试集损失和准确率到 TensorBoard
    writer.add_scalar('Test MSE Loss', mse_loss, num_epochs)
    writer.add_scalar('Test Accuracy (%)', accuracy_percentage, num_epochs)

# 可视化部分样本的真实值和预测值
num_samples_to_plot = min(5, len(true_values))  # 最多绘制 5 个样本
fig, axes = plt.subplots(num_samples_to_plot, 1, figsize=(10, num_samples_to_plot * 3), squeeze=False)
for i in range(num_samples_to_plot):
    axes[i, 0].plot(true_values[i], label="True Value")
    axes[i, 0].plot(predicted_values[i], label="Predicted Value", linestyle='--')
    axes[i, 0].set_title(f"Sample {i + 1}")
    axes[i, 0].set_xlabel("Time Step")
    axes[i, 0].set_ylabel("Value")
    axes[i, 0].legend()

plt.tight_layout()
plt.show()

# 关闭 TensorBoard writer
writer.close()
